import os
import PIL.Image as Image
import numpy as np
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
from utils.config import *
from utils.io_helper import read

# root = 'G:/SceneFlow' # Windows
# root = '/home/ubuntu/zhouyiqing/SceneFlow' # Ubuntu

# SceneFlow数据集
class SceneFlow(data.Dataset):
    def __init__(self, root, transform):
        self.root = root
        self.transform = transform

        left_imgs = list()
        right_imgs = list()
        disp_imgs = list()

        driving_left_imgs, driving_right_imgs, driving_disp_imgs = self._add_driving()
        monkaa_left_imgs, monkaa_right_imgs, monkaa_disp_imgs = self._add_monkaa()
        flyingthings3d_left_imgs, flyingthings3d_right_imgs, flyingthings3d_disp_imgs = self._add_flyingthings3d()

        left_imgs = driving_left_imgs + monkaa_left_imgs + flyingthings3d_left_imgs
        right_imgs = driving_right_imgs + monkaa_right_imgs + flyingthings3d_right_imgs
        disp_imgs = driving_disp_imgs + monkaa_disp_imgs + flyingthings3d_disp_imgs

        self.left_imgs = left_imgs
        self.right_imgs = right_imgs
        self.disp_imgs = disp_imgs

    # 获取Driving场景数据
    def _add_driving(self):
        root_driving = os.path.join(self.root, 'Driving')

        img_fmt = '{:04}.png'
        disp_fmt = '{:04}.pfm'
        driving_left_imgs = list()
        driving_right_imgs = list()
        driving_disp_imgs = list()

        driving_path = os.path.join(root_driving, 'frames_finalpass')
        disp_driving_path = os.path.join(root_driving, 'disparity')

        driving_towards = list() # 不同焦距朝向，后前后前，4个
        driving_towards.append(os.path.join(driving_path, '15mm_focallength/scene_backwards'))
        driving_towards.append(os.path.join(driving_path, '15mm_focallength/scene_forwards'))
        driving_towards.append(os.path.join(driving_path, '35mm_focallength/scene_backwards'))
        driving_towards.append(os.path.join(driving_path, '35mm_focallength/scene_forwards'))

        # 拼接图片文件名，并加入文件名list
        for towards in driving_towards:
            for i in range(1, 301):
                driving_left_imgs.append(os.path.join(towards, 'fast/left', img_fmt.format(i)))
                driving_right_imgs.append(os.path.join(towards, 'fast/right', img_fmt.format(i)))

            for j in range(1, 801):
                driving_left_imgs.append(os.path.join(towards, 'slow/left', img_fmt.format(j)))
                driving_right_imgs.append(os.path.join(towards, 'slow/right', img_fmt.format(j)))

        disp_driving_towards = list() # 不同焦距朝向，后前后前，4个
        disp_driving_towards.append(os.path.join(disp_driving_path, '15mm_focallength/scene_backwards'))
        disp_driving_towards.append(os.path.join(disp_driving_path, '15mm_focallength/scene_forwards'))
        disp_driving_towards.append(os.path.join(disp_driving_path, '35mm_focallength/scene_backwards'))
        disp_driving_towards.append(os.path.join(disp_driving_path, '35mm_focallength/scene_forwards'))

        # 拼接视差文件名，并加入文件名list
        for i in disp_driving_towards:
            for j in range(1, 301):
                driving_disp_imgs.append(os.path.join(i, 'fast/left', disp_fmt.format(j)))

            for k in range(1, 801):
                driving_disp_imgs.append(os.path.join(i, 'slow/left', disp_fmt.format(k)))

        return driving_left_imgs, driving_right_imgs, driving_disp_imgs

    # 获取Monkaa场景数据
    def _add_monkaa(self):
        root_monkaa = os.path.join(self.root, 'Monkaa')

        monkaa_left_imgs = list()
        monkaa_right_imgs = list()
        monkaa_disp_imgs = list()

        monkaa_path = os.path.join(root_monkaa, 'frames_finalpass')
        disp_monkaa_path = os.path.join(root_monkaa, 'disparity')

        monkaa_scenes = os.listdir(monkaa_path) # 获取所有场景文件夹

        # 拼接图片文件名，并加入文件名list
        for scene in monkaa_scenes:
            scene_left = os.path.join(monkaa_path, scene, 'left')
            monkaa_left_imgs.extend(os.path.join(scene_left, img) for img in os.listdir(scene_left))

            scene_right = os.path.join(monkaa_path, scene, 'right')
            monkaa_right_imgs.extend(os.path.join(scene_right, img) for img in os.listdir(scene_right))

            disp_scene = os.path.join(disp_monkaa_path, scene, 'left')
            monkaa_disp_imgs.extend(os.path.join(disp_scene, img) for img in os.listdir(disp_scene))

        return monkaa_left_imgs, monkaa_right_imgs, monkaa_disp_imgs

    # 获取FlyingThings3D场景数据
    def _add_flyingthings3d(self):
        root_flyingthings3d = os.path.join(self.root, 'FlyingThings3D')

        flyingthings3d_left_imgs = list()
        flyingthings3d_right_imgs = list()
        flyingthings3d_disp_imgs = list()

        flyingthings3d_path = os.path.join(root_flyingthings3d, 'frames_finalpass')
        disp_flyingthings3d_path = os.path.join(root_flyingthings3d, 'disparity')

        for mode in ['TRAIN']: # ['TEST', 'TRAIN']
            for abc in ['A', 'B', 'C']:
                img_path_abc = os.path.join(flyingthings3d_path, mode, abc)
                disp_path_abc = os.path.join(disp_flyingthings3d_path, mode, abc)
                flyingthings3d_folders = os.listdir(img_path_abc) # 获取所有图像文件夹

                for folder in flyingthings3d_folders:
                    folder_left = os.path.join(img_path_abc, folder, 'left')
                    flyingthings3d_left_imgs.extend(os.path.join(folder_left, img) for img in os.listdir(folder_left))

                    folder_right = os.path.join(img_path_abc, folder, 'right')
                    flyingthings3d_right_imgs.extend(os.path.join(folder_right, img) for img in os.listdir(folder_right))

                    folder_disp = os.path.join(disp_path_abc, folder, 'left')
                    flyingthings3d_disp_imgs.extend(os.path.join(folder_disp, img) for img in os.listdir(folder_disp))

        return flyingthings3d_left_imgs, flyingthings3d_right_imgs, flyingthings3d_disp_imgs

    def __getitem__(self, index):
        # 读取RGB图像
        left_img = Image.open(self.left_imgs[index]).convert('RGB')
        right_img = Image.open(self.right_imgs[index]).convert('RGB')
        if system == 'Windows': img_name = self.left_imgs[index].split('\\')[-1]
        elif system == 'Ubuntu': img_name = self.left_imgs[index].split('/')[-1]

        # 读取视差pfm，并转换为灰度图像
        disp_img = read(self.disp_imgs[index])
        disp_img = Image.fromarray(disp_img.astype(np.uint8), mode='L')

        # 图像预处理
        if self.transform:
            seed = torch.random.seed()
            torch.random.manual_seed(seed)
            left_img = self.transform(left_img)
            torch.random.manual_seed(seed)
            right_img = self.transform(right_img)

            disp_transform = transforms.Compose([transforms.ToTensor()])
            torch.random.manual_seed(seed)
            disp_img = disp_transform(disp_img)[0, :] * 255.0 # 灰度图ToTensor自动归一化到[0, 1]，还原到[0, 255]

            # 张量转图片显示
            # transforms.ToPILImage()(left_img).convert('RGB').show()
            # transforms.ToPILImage()(right_img).convert('RGB').show()
            # transforms.ToPILImage()(disp_img).convert('L').show()

            # 数据保存到字典
            data = {'left': left_img, 'right': right_img, 'gt': disp_img, 'name': img_name}

        return data

    def __len__(self):
        return len(self.left_imgs)
